import datetime
import json
import multiprocessing
import os
import random
import uuid
from collections import defaultdict
from copy import copy

import cv2
import numpy as np

from Ling.common.utils.OcrUtils import get_lines_from_file, \
    get_random_items, filter_line_in, get_dict_from_file, get_file_mapping, sort_mapping_by_value, labels_split
from text_renderer.config import get_cfg
from text_renderer.render import Render

gl_total = 0
gl_char_img_num = defaultdict(int)
out_path = 'D:/ocr/ocr_resources/reader/train_data/train_data_v3'
font_root = 'D:/ocr/ocr_resources/reader/fonts/'
save_path = 'D:/ocr/ocr_resources/reader/train_data/process_char_enum_by_font/'
char_root = 'D:/ocr/ocr_resources/reader/dict_10_no_dot/'
pre_char_texts_mapping_path = 'D:/ocr/ocr_resources/reader/all_words/pre_json_char_texts_mapping.txt'

lock = multiprocessing.Lock()


class EnumRender:
    def __init__(
            self,
            num_image=None,
            save_dir=None,
            texts=None,
            font_path=None
    ):
        generator_cfgs = get_cfg('config/enum.py')
        generator_cfg = generator_cfgs[0]
        render_cfg = generator_cfg.render_cfg
        corpus = render_cfg.corpus
        corpu = corpus[0]
        corpus_cfg = corpu.cfg
        corpus_cfg.text_paths = []
        self.generator_cfgs = generator_cfgs

        if num_image:
            self.set_num_image(num_image)
        if save_dir:
            self.set_save_dir(save_dir)
        if texts:
            self.set_texts(texts)
        if font_path:
            self.set_font_path(font_path)

    def set_texts(self, texts):
        corpu = self.generator_cfgs[0].render_cfg.corpus[0]
        corpu.texts = texts

    def set_font_path(self, font_path):
        if font_path is not None:
            corpu = self.generator_cfgs[0].render_cfg.corpus[0]
            corpu.cfg.font_list_file = font_path
            corpu.font_manager.load_font(font_path, None)

    def set_num_image(self, num_image):
        generator_cfg = self.generator_cfgs[0]
        generator_cfg.num_image = num_image

    def set_save_dir(self, save_dir):
        generator_cfg = self.generator_cfgs[0]
        generator_cfg.save_dir = save_dir

    def set_img_pre_name(self, img_pre_name):
        self.img_pre_name = img_pre_name

    def write_datas(self, datas):
        save_dir = self.generator_cfgs[0].save_dir
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)
        _img_dir = os.path.join(save_dir, "images")
        if not os.path.exists(_img_dir):
            os.makedirs(_img_dir)
        label_path = os.path.join(save_dir, 'labels.txt')

        lock.acquire()
        with open(label_path, "a", encoding="utf-8") as f:
            for pos, data in enumerate(datas):
                name_uuid = str(uuid.uuid1())
                pre = self.img_pre_name
                name = pre + '_' + name_uuid[0:8] + name_uuid[19:23]
                label = data['label']
                # 换行bug
                label = label.replace('\r', '').replace('\n', '').replace('\t', '')
                image = data['image']
                img_path = os.path.join(_img_dir, name + ".jpg")
                cv2.imwrite(img_path, image, [int(cv2.IMWRITE_JPEG_QUALITY), 95])
                height, width = image.shape[:2]
                line = name + '\t' + label + '\t' + str(height) + '\t' + str(width)
                f.write(line + '\n')
        lock.release()

    def process_img(self):
        save_dir = self.generator_cfgs[0].save_dir
        if self.generator_cfgs[0].render_cfg.corpus[0].texts is None: raise 'texts不能为空'
        if self.generator_cfgs[0].num_image is None: raise 'num_image不能为空'
        if save_dir is None: raise 'save_dir不能为空'
        if self.img_pre_name is None: raise 'img_pre_name图片名称前缀不能为空'
        if self.generator_cfgs[0].render_cfg.corpus[0].cfg.font_list_file is None: raise 'font_path不能为空'
        for generator_cfg in self.generator_cfgs:
            np.random.seed()
            render = Render(generator_cfg.render_cfg)
            datas = []
            for _ in range(generator_cfg.num_image):
                data = render()
                if data is not None:
                    datas.append({"image": data[0], "label": data[1]})
        self.write_datas(datas)


root = 'D:/workspace/python/train_data/'
dictRoot = root + 'train_data_configs/words/'
reader_root = 'D:/ocr/ocr_resources/reader/'


def process_char_lines(dict_in, char, line_num, reader_dict_root):
    print('begin load char:' + char)
    lines = get_lines_from_file(reader_dict_root + char + '.txt')
    if len(lines) > 0:
        get_line_num = int(line_num * 1.5)
        # 取得字符串
        lines = get_random_items(lines, get_line_num)
        if dict_in is None:
            return lines
        else:
            return filter_char_lines(lines, dict_in, char)
    else:
        return []


def filter_char_lines(lines, dict_in, char):
    # 保证字符串的文字都能生成
    lines = filter_line_in(lines, dict_in)
    # 取得的是与char有关的字符串
    lines_filter = []
    for line in lines:
        if line.find(char) >= 0:
            lines_filter.append(line)
    return lines_filter


def load_char_texts_mapping(pre_char_texts_mapping_path):
    with open(pre_char_texts_mapping_path, 'r', encoding='utf-8') as f:
        char_texts_mapping = json.load(f)
    return char_texts_mapping


def load_font_render_mapping(font_root, save_path):
    fonts = [f for f in os.listdir(font_root) if f.lower().endswith('ttc') or f.lower().endswith('ttf')]
    font_render_mapping = defaultdict()
    for pos, font in enumerate(fonts):
        if pos < 9999:
            print('begin load font:{} {} of {}'.format(font, pos, len(fonts)))
            enum_render = EnumRender(save_dir=save_path)
            enum_render.set_font_path(font_root + font)
            font_render_mapping[font] = enum_render
    return font_render_mapping


def load_font_char_mapping(font_root):
    fonts = [f for f in os.listdir(font_root) if f.lower().endswith('ttc') or f.lower().endswith('ttf')]
    # font支持的字符映射
    font_char_mapping = defaultdict()
    # 预加载数据
    for pos, font in enumerate(fonts):
        print('begin load font:{} {} of {}'.format(font, pos, len(fonts)))
        font_dict_in = get_dict_from_file(font_root + font + '.txt')
        font_char_mapping[font] = font_dict_in
    return font_char_mapping


def get_char_texts_mapping(char_num_path, max_num):
    char_img_num = defaultdict(int)
    char_texts_mapping = defaultdict(list)
    mapping = get_file_mapping(char_num_path, '\t')
    mapping = sort_mapping_by_value(mapping)
    for pos, char in enumerate(mapping.keys()):
        char_num = char_img_num[char]
        if char_num <= max_num:
            texts = process_char_lines(None, char, 200, char_root)
            for text in texts:
                for c in text:
                    char_img_num[c] = char_img_num[c] + 1
            char_texts_mapping[char] = texts
    char_img_num_sort = sort_mapping_by_value(char_img_num)
    return char_texts_mapping, char_img_num_sort


def pre_json_char_texts_mapping(char_num_path, out_path, max_num):
    char_texts_mapping, char_img_num_sort = get_char_texts_mapping(char_num_path, max_num)
    with open(out_path, 'w', encoding='utf-8') as f:
        json.dump(char_texts_mapping, f, indent=4, ensure_ascii=False, sort_keys=True)
    with open(out_path.replace('.txt', '') + '_char_num.txt', 'w', encoding='utf-8') as f:
        for key in char_img_num_sort:
            f.write(key + '\t' + str(char_img_num_sort[key]) + '\n')


def render_font_img(font_render_mapping, font_char_mapping, char_texts_mapping, font):
    global gl_total
    global char_img_num
    enum_render = font_render_mapping[font]
    font_dict_in = font_char_mapping[font]

    with open('D:/ocr/ocr_resources/reader/train_data/log.txt', 'a', encoding='utf-8') as log:
        for char in font_dict_in:
            texts = char_texts_mapping[char]
            if len(texts) > 0:
                # 随机抽取200次,不足200的就重复
                ran = np.random.randint(0, len(texts), 200)
                text = texts[ran]
                render_texts = filter_char_lines([text], font_dict_in, char)
                if len(render_texts) > 0:
                    gl_total = gl_total + 1
                    info = 'font:{} char:{} text:{} total:{}'.format(font, char, text, gl_total)
                    print(info)
                    log.write(info + '\n')
                    enum_render.set_num_image(1)
                    enum_render.set_texts(render_texts)
                    enum_render.process_img()
                    for c in render_texts[0]:
                        char_img_num[c] = char_img_num[c] + 1


def get_rand_support_font(text, font_char_mapping, font_render_mapping):
    fonts = list(font_render_mapping.keys())
    font = random.choice(fonts)
    char_dict_in = font_char_mapping[font]
    if is_text_in_dict(text, char_dict_in):
        return font_render_mapping[font], font
    return get_rand_support_font(text, font_char_mapping, font_render_mapping)


def is_text_in_dict(text, dict):
    for char in text:
        if char not in dict:
            return False
    return True


def t_error(save_path):
    text = '㐶'
    files = [f for f in os.listdir(font_root) if f.lower().endswith('ttc') or f.lower().endswith('ttf')]
    # font = '方正悠黑系列_511M.ttf'
    # font = 'Deng.ttf'
    for f in files:
        try:
            enum_render = EnumRender(save_dir=save_path)
            enum_render.set_font_path(font_root + f)
            enum_render.set_num_image(1)
            enum_render.set_texts([text])
            enum_render.process_img()
        except:
            print(f)


def run_by_round(font_char_mapping, char_texts_mapping, font_render_mapping, img_pre_name):
    # 每轮每个字符至少出现一次,当出现次数大于10,不再为其专门读取texts生成图片
    # 10可能有点低,那些常用字符的关联字符会出现过拟合
    # 此方法执行200次,那么,每个字符最少出现200次
    # 有此方法后就可以多线程了
    global gl_total
    round_char_img_num = defaultdict(int)
    with open('D:/ocr/ocr_resources/reader/train_data/log.txt', 'a', encoding='utf-8') as log:
        for char in char_texts_mapping:
            char_num = round_char_img_num[char]
            if char_num <= 10:
                text = random.choice(char_texts_mapping[char])
                enum_render, font = get_rand_support_font(text, font_char_mapping, font_render_mapping)
                info = 'total:{} char:{} in text:{} with font:{}'.format(gl_total, char, text, font)
                # print(info)
                log.write(info + '\n')
                enum_render.set_num_image(1)
                enum_render.set_texts([text])
                enum_render.set_img_pre_name(img_pre_name)
                enum_render.process_img()
                for c in text:
                    round_char_img_num[c] = round_char_img_num[c] + 1
                gl_total = gl_total + 1


def round_by_pool(font_root, save_path, pre_char_texts_mapping_path):
    pool = multiprocessing.Pool(processes=8)
    font_char_mapping = load_font_char_mapping(font_root)
    char_texts_mapping = load_char_texts_mapping(pre_char_texts_mapping_path)
    font_render_mapping = load_font_render_mapping(font_root, save_path)
    for i in range(0, 200):
        font_render_mapping_ = defaultdict()
        for font in font_render_mapping:
            # print('begin copy font:' + font)
            font_render_mapping_[font] = copy(font_render_mapping[font])
        pool.apply_async(run_by_round,
                         (font_char_mapping, char_texts_mapping, font_render_mapping_, str(i)))
        # run_by_round(font_char_mapping, char_texts_mapping, font_render_mapping_)
    pool.close()
    pool.join()


if __name__ == "__main__":
    starttime = datetime.datetime.now()

    # dict_in = get_char_in()

    # save_path = 'D:/ocr/ocr_resources/reader/train_data/process_char_enum_by_font/'
    # enum_render = EnumRender(save_dir=save_path)
    # enum_render.set_num_image(1)
    # enum_render.set_texts(['道压迫症状最为常见表'])
    # enum_render.set_font_path('D:/ocr/ocr_resources/reader/fonts/Dengb.ttf')
    # enum_render.process_img()
    #
    # enum_render.set_num_image(3)
    # enum_render.set_texts(['参考位置为鼻咽顶黏膜'])
    # enum_render.set_font_path('D:/ocr/ocr_resources/reader/fonts/FZSJ-SHENQXBGT.TTF')
    # enum_render.process_img()
    # reader_dict_root = reader_root + 'dict_10/'
    # reader_dict_root_no_dot = reader_root + 'dict_10_no_dot/'
    # mapping = get_file_mapping(reader_root + 'all_words/char_num_sort_10.txt', '\t')
    # mapping_no_dot = get_file_mapping(reader_root + 'all_words/char_num_sort_10_no_dot.txt', '\t')

    # for font in font_render_mapping.keys():
    #     render = font_render_mapping[font]
    #     font_dict_in = font_char_mapping[font]
    #     for char in font_dict_in:
    #         texts = char_texts_mapping[char]
    #         ran = random.randint(0, len(texts) - 1)
    #         text = texts[ran]
    #         enum_render.set_num_image(1)
    #         enum_render.set_texts([text])
    #         enum_render.process_img()

    # pre_json_char_texts_mapping(reader_root + 'all_words/char_num_sort_10_no_dot.txt', pre_char_texts_mapping_path, 500)

    # for font in font_render_mapping.keys():
    #     try:
    #         enum_render = font_render_mapping[font]
    #         dict = font_char_mapping[font]
    #         if '㐶' in dict:
    #             enum_render.set_num_image(1)
    #             enum_render.set_texts(['㐶'])
    #             enum_render.process_img()
    #     except:
    #         print(font)

    # with open('D:/ocr/ocr_resources/reader/train_data/log.txt', 'a', encoding='utf-8') as log:
    #     for char in char_texts_mapping:
    #         char_num = gl_char_img_num[char]
    #         if char_num <= 500:
    #             for i in range(0, 200):
    #                 text = random.choice(char_texts_mapping[char])
    #                 enum_render, font = get_rand_support_font(text, font_char_mapping, font_render_mapping)
    #                 info = 'total:{} rand:{} char:{} in text:{} with font:{}'.format(gl_total, i, char, text, font)
    #                 print(info)
    #                 log.write(info + '\n')
    #                 enum_render.set_num_image(1)
    #                 enum_render.set_texts([text])
    #                 enum_render.process_img()
    #                 for c in text:
    #                     gl_char_img_num[c] = gl_char_img_num[c] + 1
    #                 gl_total = gl_total + 1

    # t_error(save_path)
    # ********************************v3 begin*************************#

    font_char_mapping = load_font_char_mapping(font_root)
    char_texts_mapping = load_char_texts_mapping(pre_char_texts_mapping_path)
    # char_texts_mapping, char_img_num_sort = get_char_texts_mapping(reader_root + 'all_words/char_num_sort_10_no_dot.txt', pre_char_texts_mapping_path, 500)
    # font_render_mapping = load_font_render_mapping(font_root, save_path)
    # run_by_round(font_char_mapping, char_texts_mapping, font_render_mapping, '1')

    round_by_pool(font_root, save_path, pre_char_texts_mapping_path)
    labels_split('D:/ocr/ocr_resources/reader/',
                 ['enum_by_font'], 10000)
    # ********************************v3 end*************************#
    endtime = datetime.datetime.now()
    seconds = (endtime - starttime).seconds
    m, s = divmod(seconds, 60)
    h, m = divmod(m, 60)
    print("%d:%02d:%02d" % (h, m, s))
